# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 14:17 
@Author  : ZhangShenao 
@File    : 1.通用RAG系统处理流程.py 
@Desc    : 通用RAG系统处理流程

一个通用的RAG系统的处理流程
1. loading(文档加载)
2. splitting/transform(文档切割/文档转换)
3. embedding/save(文档嵌入/向量存储)
4. retrieving(文档检索)

"""

import dotenv
from langchain_community.chat_models.moonshot import MoonshotChat
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载环境变量
dotenv.load_dotenv()

MoonshotChat

# 1. loading(文档加载)
loader = UnstructuredFileLoader('./docs/骆驼祥子.txt')
origin_docs = loader.load()

# 2. 文档切割
separators = [
    '\n\n',
    '\n',
    r'。|！|？',
    r'\.\s|\!\s|\?\s',  # 英文标点符号后面通常需要加空格
    r'；|;\s',
    r'，|,\s',
    ' ',
    ''
]
splitter = RecursiveCharacterTextSplitter(
    separators=separators,  # 指定自定义分隔符列表
    chunk_size=20,  # 指定文档分块大小
    chunk_overlap=5,  # 指定文档块之间的重叠部分大小
    add_start_index=True,  # 添加文档块的起始索引
    is_separator_regex=True,  # 指定分隔符为正则表达式形式
)
split_docs = splitter.split_documents(origin_docs)

# 3. Embedding存储
embeddings = ZhipuAIEmbeddings(model='embedding-2')
vector_store = PineconeVectorStore(
    index_name='dataset',  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace='llm-ops',  # 指定Namespace
    text_key='text',  # 指定文本属性key
)
vector_store.add_documents(split_docs)

# 4. Retrieving文档检索
retrieve_result = vector_store.similarity_search_with_relevance_scores(
    query='骆驼祥子与刘四爷是什么关系？',
    k=2,
    score_threshold=0.5,  # 设置相关性分数阈值,只有相关性超过阈值的文档才会被检索出来
)
print(retrieve_result)
